import argparse

import os
from PIL import Image


def extract_prefix_from_filename(filename):
    parts = filename.split("_")
    prefix = parts[0]
    return prefix


def get_unique_prefix_from_folder(input_folder):
    prefix_set = set()

    for filename in os.listdir(input_folder):
        if filename.endswith(".png"):
            prefix = extract_prefix_from_filename(filename)
            prefix_set.add(prefix)

    if len(prefix_set) != 1:
        raise ValueError("Multiple unique prefixes found in the folder.")

    return prefix_set.pop()


def calculate_original_size(input_folder):
    grid_rows = 0
    grid_cols = 0
    tile_size = 0

    for filename in os.listdir(input_folder):
        if filename.endswith(".png"):
            file_parts = filename.split("_")
            row_col_part = file_parts[1].split(".")[0].split("-")  # "tr1-tc1"
            row = int(row_col_part[0].split("tr")[-1])  # 提取出行
            col = int(row_col_part[1].split("tc")[-1])  # 提取出列

            if row > grid_rows:
                grid_rows = row
            if col > grid_cols:
                grid_cols = col

            image_path = os.path.join(input_folder, filename)
            image = Image.open(image_path)
            tile_size = image.width

    original_width = grid_cols * tile_size
    original_height = grid_rows * tile_size
    original_size = (original_width, original_height)
    return original_size


def downsample_and_merge_images(input_folder, output_folder, downsample_ratio, original_size, tile_size):
    prefix = get_unique_prefix_from_folder(input_folder)  # 获取统一前缀
    max_row = int(original_size[1]*downsample_ratio / tile_size)  # 计算最大行数
    max_col = int(original_size[0]*downsample_ratio / tile_size)  # 计算最大列数

    for cur_row in range(max_row):
        for cur_col in range(max_col):
            merged_image = Image.new("RGB", (tile_size, tile_size))

            for row_offset in range(2):
                for col_offset in range(2):
                    filename = "{}_tr{}-tc{}.png".format(prefix, cur_row * 2 + row_offset + 1,
                                                         cur_col * 2 + col_offset + 1)
                    image_path = os.path.join(input_folder, filename)

                    if os.path.exists(image_path):
                        image = Image.open(image_path)

                        # 下采样
                        new_width = int(image.width * downsample_ratio)
                        new_height = int(image.height * downsample_ratio)
                        image = image.resize((new_width, new_height))

                        # 计算小图像在合并图像中的位置
                        x = col_offset * image.width
                        y = row_offset * image.height

                        # 将小图像放置到合并图像中的对应位置
                        merged_image.paste(image, (x, y))

            # 保存合并后的图像
            output_filename = "{}_tr{}-tc{}.png".format(prefix, cur_row + 1, cur_col + 1)
            output_path = os.path.join(output_folder, output_filename)
            os.makedirs(output_folder, exist_ok=True)  # 创建输出目录
            merged_image.save(output_path)


def main():
    parser = argparse.ArgumentParser(description='Downsample the spliced and rendered image')
    parser.add_argument("--in_dir", type=str,
                        default="E:\\wafer52\\11866_16nm",
                        help="image dir")
    parser.add_argument("--out_dir", type=str,
                        default="E:\\wafer52\\11866_32nm",
                        help="output dir")
    parser.add_argument("--down_sample_ratio", type=float, default=0.5,
                        help="set the down_sample_ratio")
    parser.add_argument("--tile_size", type=int, default=4096,
                        help="the size (square side) of each tile, -1 means the whole image. (Set the value to a multiple of 128(for neuroglancer display) or -1)")
    args = parser.parse_args()
    original_size = calculate_original_size(args.in_dir)
    downsample_and_merge_images(args.in_dir, args.out_dir, args.down_sample_ratio, original_size, args.tile_size)


if __name__ == '__main__':
    main()
